import random
import numpy as np
import time

from src.generator_utils import get_random_seed_tensor, exec_method_by_index
from src.evaluator import has_precision_problem_occurred
import tensorflow as tf

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"


def init_group(init_num, shape, operator, pattern, b, r):
    group_tf = []
    group_torch = []
    group_mnn = []
    while len(group_tf) < init_num or len(group_torch) < init_num or len(group_mnn) < init_num:
        tensor_seed = get_random_seed_tensor(shape)
        success = False
        for index in range(42):
            new_tensor = exec_method_by_index(tensor_seed, index)
            if not isinstance(new_tensor, str):
                detect_res = has_precision_problem_occurred(new_tensor, operator, pattern, b, r)
                if detect_res[0] and len(group_tf) < init_num:
                    group_tf.append(new_tensor)
                if detect_res[1] and len(group_torch) < init_num:
                    group_torch.append(new_tensor)
                if detect_res[2] and len(group_mnn) < init_num:
                    group_mnn.append(new_tensor)
                if len(group_tf) == init_num and len(group_torch) == init_num and len(group_mnn) == init_num:
                    success = True
                    break
        if success:
            break
    return group_tf, group_torch, group_mnn


def extend_group(group, framework, final_num,  operator, pattern, b, r, shape):
    framework_lst = ['tf', 'torch', 'mnn']
    new_set = []
    for member in group:
        new_tensor = find_neighbor(member)
        t = 0
        while t < 30:
            detect_res = has_precision_problem_occurred(new_tensor, operator, pattern, b, r)
            framework_index = framework_lst.index(framework)
            if detect_res[framework_index]:
                new_set.append(new_tensor)
                break
            t += 1
            new_tensor = find_neighbor(member)

        if len(group) + len(new_set) >= final_num:
            break
    # 将新生成的member加入group
    new_set.extend(group)
    print('current size: {0}'.format(len(new_set)))
    return new_set


def find_neighbor(tensor_seed):
    rand_seed = random.randint(0, 6)
    if rand_seed == 0:
        new_tensor = np.random.normal(tensor_seed, 1/3)
    elif rand_seed == 1:
        new_tensor = np.random.normal(tensor_seed, 0.5/3)
    elif rand_seed == 2:
        new_tensor = np.random.normal(tensor_seed, 0.1/3)
    elif rand_seed == 3:
        new_tensor = np.random.normal(tensor_seed, 0.05/3)
    elif rand_seed == 4:
        new_tensor = np.random.normal(tensor_seed, 0.01/3)
    elif rand_seed == 5:
        new_tensor = np.random.normal(tensor_seed, 0.005/3)
    elif rand_seed == 6:
        new_tensor = np.random.normal(tensor_seed, 0.001/3)
    return new_tensor


if __name__ == '__main__':
    operator = 'conv2d'
    pattern = 'MRE'
    b = 0.003
    r = 0.001
    shape = (1, 1, 28, 28)

    # # batch_normalization层需要开启训练状态 使用样本均值和方差而不是默认均值或方差
    # tf.keras.backend.set_learning_phase(1)

    tf_group, torch_group, mnn_group = init_group(10, shape, operator, pattern, b, r)
    print(len(tf_group))
    print(len(torch_group))
    print(len(mnn_group))
    # 首先tf
    c = 1
    start_t = time.time()
    while len(tf_group) < 15000:
        print('iteration {0}...'.format(c))
        tf_group = extend_group(tf_group, 'tf', 15000, operator, pattern, b, r, shape)
        c += 1

    end_t = time.time()
    tf_duration = end_t - start_t

    # torch
    c = 1
    start_t = time.time()
    while len(torch_group) < 15000:
        print('iteration {0}...'.format(c))
        torch_group = extend_group(torch_group, 'torch', 15000, operator, pattern, b, r, shape)
        c += 1

    end_t = time.time()
    torch_duration = end_t - start_t

    # mnn
    c = 1
    start_t = time.time()
    while len(mnn_group) < 15000:
        print('iteration {0}...'.format(c))
        mnn_group = extend_group(mnn_group, 'mnn', 15000, operator, pattern, b, r, shape)
        c += 1

    end_t = time.time()
    mnn_duration = end_t - start_t

    print('successfully generated {0} samples for tf!'.format(len(tf_group)))
    print('successfully generated {0} samples for torch!'.format(len(torch_group)))
    print('successfully generated {0} samples for mnn!'.format(len(mnn_group)))

    print('during time {0} --tf'.format(tf_duration))
    print('during time {0} --tf'.format(torch_duration))
    print('during time {0} --tf'.format(mnn_duration))